# -*- coding: utf-8 -*-
"""
Created on Mon Mar  4 19:31:20 2019

@author: william

Email: hua_yan_tsn@163.com
"""
import numpy as np
from WordEmbeddings import laptop_model, restaurant_model
# the directory of laptop data and restaurant data
laptop_dir = '../DataSet/Aspect-Sentiment-Analysis-master/laptop_data/'
restaurant_dir = '../DataSet/Aspect-Sentiment-Analysis-master/restaurant_data/'
# the path of laptop's and restaurant's data file whose content is reflected from original words into index integer
laptop_corpus = laptop_dir + 'train_text_index.txt'
laptop_test_corpus = laptop_dir + 'test_text_index.txt'
restaurant_corpus = restaurant_dir + 'train_text_index.txt'
restaurant_test_corpus = restaurant_dir + 'test_text_index.txt'
# the path of laptop's label file and restaurant's label file.
laptop_label = laptop_dir + 'train_aspects_label_index.txt'
laptop_test_label = laptop_dir + 'test_aspects_label_index.txt'
restaurant_label = restaurant_dir + 'train_aspects_label.txt'
restaurant_test_label = restaurant_dir + 'test_aspects_label.txt'
# the path of file which contains the aspect of correspond sentence
laptop_aspect_file = laptop_dir + 'train_aspects_text_index.txt'
laptop_test_aspect_file = laptop_dir + 'test_aspects_text_index.txt'
restaurant_aspect_file = restaurant_dir+'train_aspect_text_index.txt'
restaurant_test_aspect_file = restaurant_dir+'test_aspect_text_index.txt'
# the path of file which is the word index reflecting table
laptop_word_index = laptop_dir + 'word_index.txt'
restaurant_word_index = restaurant_dir + 'word_index.txt'
def getWordEmbeddings(inputs, max_len = 0):
    input_list = inputs.split()
    word_embeddings = np.zeros((max_len, 300))
    print('word embeddings shape is ', word_embeddings.shape)
    for i in range(len(input_list)):
        try:
            word_embeddings[i] = laptop_model.get_vector(input_list[i])
        except:
            try:
                word_embeddings[i] = restaurant_model.get_vector(input_list[i])
            except:
                word_embeddings[i] = np.zeros(laptop_model.get_vector('much').shape)
    return word_embeddings

def loadWordIndex(source):
    if source == 'laptop':
        filepath = laptop_word_index
    elif source == 'restaurant':
        filepath = restaurant_word_index
    else:
        return {}
    with open(filepath, 'r') as file:
        data = file.readlines()
        dicts = {}
        for line in data:
            line = line.strip('\n').strip()
            word, index = line.split()
            dicts.setdefault(index, word)
        return dicts

def loadWordsIds(max_len=82):
    dicts = {}
    for position in range(-1 * (max_len - 1), max_len):
        dicts.setdefault(position, position + max_len)
    return dicts

def transfer(word2vec, word_index):
    """
    :param word2vec:  word2vec的模型对象
    :param word_index: key 为 index， value为word的字典对象
    :return: 存放matrix的文件名
    """
    dir = laptop_dir  + 'embedding_matrix.bin'
    import pickle
    import os
    if (os.path.exists(dir)):
        with open(dir, 'rb') as file:
            embeddings = pickle.load(file)
            return embeddings
    else:
        embeddings = []
        for key in word_index:
            # TODO: 此处的词向量有点问题，需要修正
            try:
                vector = word2vec.get_vector(word_index[key])
            except:
                vector = np.zeros((300, ))
            embeddings.append(vector)
        embeddings = np.array(embeddings)
        with open(dir, 'wb') as file:
            pickle.dump(embeddings, file)
            return embeddings

"""
加载数据的函数
"""
def loadData(source, type='train'):
    """
    :param source: 加载源文件的源
    :param type: 'train' | 'test' 表示是训练集或者是测试集
    :return: 返回文件内容
    """
    if (type == 'train'):
        if source == 'laptop':
            file = open(laptop_corpus, 'r', encoding='utf8')
        elif source == 'restaurant':
            file = open(restaurant_corpus, 'r', encoding = 'utf8')
        else:
            print('No such file')
            raise ValueError
    elif(type == 'test'):
        if source == 'laptop':
            file = open(laptop_test_corpus, 'r', encoding='utf8')
        elif source == 'restaurant':
            file = open(restaurant_test_corpus, 'r', encoding='utf8')
        else:
            print('No such file')
            raise ValueError
    data = file.readlines()
    file.close()
    return data

def loadAspectToken(resource, type='train'):
    """
    :param source: 加载源文件的源
    :param type: 'train' | 'test' 表示是训练集或者是测试集
    :return: 返回文件内容
    """
    if (type == 'train'):
        if (resource == 'laptop'):
            filepath = laptop_aspect_file
        elif (resource == 'restaurant'):
            filepath = restaurant_aspect_file
        else:
            return []
    elif (type == 'test'):
        if (resource == 'laptop'):
            filepath = laptop_test_aspect_file
        elif (resource == 'restaurant'):
            filepath = restaurant_test_aspect_file
        else:
            return []
    with open(filepath, 'r', encoding='utf8') as file:
        data = file.readlines()
        return data

def loadLabel(source, type='train'):
    """
    :param source:
    :return:
    """
    if (type == 'train'):
        if source == 'laptop':
            file = open(laptop_label, 'r', encoding='utf8')
        elif source == 'restaurant':
            file = open(restaurant_label, 'r', encoding = 'utf8')
        else:
            print('No such file')
            raise ValueError
    elif (type == 'test'):
        if source == 'laptop':
            file = open(laptop_test_label, 'r', encoding='utf8')
        elif source == 'restaurant':
            file = open(restaurant_test_label, 'r', encoding = 'utf8')
        else:
            print('No such file')
            raise ValueError
    data = file.readlines()
    file.close()
    label = np.zeros((len(data), 4))
    for index in range(len(data)):
        label[index][int(data[index])] = 1
    """
    the weight matrix, the first dimension of W denotes the class -1, second dimension denotes
            the class 1, and third dimension denotes the class 0 and the fourth dimension denotes the class 3
            class -1 : negative
            class 1: positive
            class 0: nerual
            class 3: conflicts
    """
    return label

def loadInitPositionEmbeddings(max_len=100, vector_dimension=20):
    import pickle
    import os
    filename = 'initPositionEmbeddings.bin'
    if (os.path.exists(filename)):
        print('====initial position embeddings file exists, directly load====')
        with open(filename, 'rb') as file:
            embeddings = pickle.load(file)
    else:
        print('====initial position embeddings file not exists, create and save====')
        length = max_len * 2
        embeddings = np.random.rand(length, vector_dimension)
        with open(filename, 'wb') as file:
            pickle.dump(embeddings, file)
    return embeddings

